import os
import base64
from datetime import datetime
from google import genai
from google.genai import types
from google.cloud import storage
from google.genai.types import HttpOptions, Part
from prompts import prompt_map, language_map
from typing import List, Dict, Optional, Any

class GoogleSpeechTranscriber:
    def __init__(self, project_id, location="global", model="gemini-2.5-flash", bucket="cloud-ai-platform-1a06e8cb-b424-44c0-a95e-c28718a96ba3"):
        self.client = genai.Client(vertexai=True, project=project_id, location=location)
        self.model = model
        self.bucket = bucket

    def upload_to_gcs(self, source_file_path):
        print("Uploading to GCS...")
        storage_client = storage.Client()
        bucket = storage_client.bucket(self.bucket)
        filename = os.path.basename(source_file_path)
        destination_blob_name = f"audio_uploads/{filename}"
        blob = bucket.blob(destination_blob_name)
        blob.upload_from_filename(source_file_path)
        gcs_uri = f"gs://{self.bucket}/{destination_blob_name}"
        print(f"Uploaded to {gcs_uri}")
        return gcs_uri

    def delete_from_gcs(self, source_file_path):
        print("Deleting GCS file...")
        filename = os.path.basename(source_file_path)
        destination_blob_name = f"audio_uploads/{filename}"
        storage_client = storage.Client()
        bucket = storage_client.bucket(self.bucket)
        blob = bucket.blob(destination_blob_name)
        blob.delete()
        print(f"Deleted gs://{self.bucket}/{destination_blob_name}")

    def transcribe_audio(self, audio_file_path, language):
        if not os.path.exists(audio_file_path):
            return {"error": f"File not found: {audio_file_path}"}

        # gcs_uri = self.upload_to_gcs(audio_file_path)
        with open(audio_file_path, "rb") as f:
            audio_bytes = f.read()

        audio_part = types.Part.from_bytes(
            data=audio_bytes,
            mime_type="audio/mpeg",  # 或根据实际情况使用 "audio/wav", "audio/webm"
        )

        # language = "Chinese"
        # egLanguage = "\"马来语\", \"泰语\""
        egLanguage = language_map[language]
        language = language.replace("_", " ")
        contents = [
            types.Content(
                role="user",
                parts=[
                    audio_part,
                    types.Part.from_text(
                       text="Task: Please identify the content of the audio file I send.\n"
                            "If the spoken language is Vietnamese, English, Thai, Malay, or Japanese (i.e. non-"+language+"), you must translate **each sentence** into "+language+".\n"
                            "**Important: Do NOT treat Vietnamese original text as the translation. You must provide proper "+language+" translations.**\n"
                            "Step-by-step instructions:\n"
                           "- If you did not receive any file, reply with: {\"error\": \"没有收到文件\"}\n"
                           "- If you received a file but it is not a recognizable audio file, reply with: {\"error\": \"无法正确识别音频\"}\n"
                           "- If the file is an audio file but contains no speech, reply with: {\"error\": \"文件中没有人再说话\"}\n"
                           "- If the file contains spoken content, detect the language and transcribe the speech.\n"
                           "- Translate each sentence into "+language+" if the language is not "+language+".\n"
                           "- If the content cannot be parsed or understood, reply with: {\"error\": \"无法正确解析\"}\n"
                           "**Output Format Requirements:**\n"
                           "1. Return the detected source language in "+language+" (e.g., "+egLanguage+").\n"
                           "2. Return a JSON object containing the transcription and translation. Each entry must have timestamps accurate to the second.\n"
                           "3. Output only the pure JSON — do not include any explanations or Markdown formatting (e.g. no ```json). Return clean JSON only, suitable for programmatic parsing.\n"
                            "4. The expected JSON format is as follows (remove newlines when returning):\n"
                            "{"
                            '  "source_language": "马来语",'
                            '  "transcriptions": ['
                            "    {"
                            '      "start_time": "00s",'
                            '      "end_time": "05s",'
                            '      "original_text": "xxxx",'
                            '      "translated_text": "xxxx"'
                            "    }"
                            "  ]"
                            "}"
                    ),
                    # 加这个：专门强调越南语翻译必须正确！
                    types.Part.from_text(
                        text="Special Reminder: If the detected language is Vietnamese, you **must** accurately translate every sentence into Chinese. Never return the original Vietnamese as the translation."
                    )
                ],
            )
        ]

        generate_content_config = types.GenerateContentConfig(
            temperature=0.2,
            top_p=0.8,
            max_output_tokens=65535,
            response_modalities=["TEXT"],
            safety_settings=[
                types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="OFF"),
                types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="OFF"),
                types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="OFF"),
                types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="OFF"),
            ]
        )

        response_text = ""
        try:
            response = self.client.models.generate_content(
                model=self.model,
                contents=contents,
                config=generate_content_config,
            )

            response_text = response.text
        except Exception as e:
            print(f"Error: {str(e)}")
            return {"error": str(e)}
        # finally:
            # self.delete_from_gcs(audio_file_path)

        return {"transcription": response_text}

    def transcribe(
        self,
        text: str,
        prompt: Optional[str] = None,
        type: Optional[str] = None,
        history: Optional[List[Dict[str, str]]] = None,
    ) -> Dict[str, Any]:
        if not text:
            return {"error": "输入文本不能为空"}

        # 修复：使用 self.prompt_map 并添加类型检查
        if prompt:
            prompt_text = prompt
        elif type and type in prompt_map:
            prompt_text = prompt_map[type]
        else:
            return {"error": "未提供有效的提示词（prompt 或 type）"}

        response_text = ""
        print_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        print(f"[{print_time}] : 开始AI")
        try:
            contents = []
            # 添加历史对话内容（如果有）
            if history:
                for item in history:
                    role = item.get("role", "")
                    message = item.get("text", "")
                    if role and message:
                        contents.append(
                            types.Content(
                                role=role,
                                parts=[types.Part.from_text(message)]
                            )
                        )

            # 构建内容
            text_input_part = types.Part.from_bytes(data=text.encode("utf-8"), mime_type="text/plain")
            prompt_part = types.Part.from_text(text=prompt_text)
            contents.append(
                types.Content(
                    role="user",
                    parts=[text_input_part, prompt_part]
                )
            )
            print_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            print(f"[{print_time}] : 构造请求配置")
            # 构造请求配置
            generate_content_config = types.GenerateContentConfig(
                temperature=0.2,
                top_p=0.8,
                max_output_tokens=8192,
                response_modalities=["TEXT"],
                safety_settings=[
                    types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="OFF"),
                    types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="OFF"),
                    types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="OFF"),
                    types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="OFF"),
                ]
            )
            print_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            print(f"[{print_time}] : Config：end")
            response = self.client.models.generate_content(
                model=self.model,
                contents=contents,
                config=generate_content_config,
            )
            response_text = response.text
            print_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            print(f"[{print_time}] : 请求google结束")
        except Exception as e:
            # 添加详细错误日志
            import traceback
            error_detail = traceback.format_exc()
            print(f"转录错误: {str(e)}\n{error_detail}")
            return {"error": f"处理失败: {str(e)}"}
        print_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        print(f"[{print_time}] : 返回数据：{response.text}")
        return {"transcription": response.text}